import numpy as np
import matplotlib.pyplot as plt

# 1. 设置中文字体和整体风格（确保负号正常显示）
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial'] # 用SimHei显示中文，Arial作为后备
plt.rcParams['axes.unicode_minus'] = False
plt.style.use('default') # 使用清晰的白色背景

# 2. 准备数据
# 横坐标：类别数量 K (从1到10，取整数)
K = np.arange(1, 11) # 生成数组 [1, 2, 3, ..., 10]

# 当所有类别等概率时，每个类的概率 P = 1/K
P_equal = 1.0 / K # 纵坐标右轴数据：概率值

# 计算最大熵（当所有类别等概率时，熵达到最大）: H_max = -Σ(1/K * log2(1/K)) = log2(K)
H_max = np.log2(K) # 纵坐标左轴数据：熵值

# 3. 创建画布和第一个坐标轴
fig, ax1 = plt.subplots(figsize=(10, 6)) # 设置图片大小

# 4. 绘制左轴数据：最大熵曲线 (红色，方形标记)
line_entropy, = ax1.plot(K, H_max, 's-', color='red', linewidth=2, markersize=8, label='Max Entropy')
ax1.set_xlabel('Number of classes', fontsize=12)
ax1.set_ylabel('Maximum Entropy', color='red', fontsize=12) # 左轴标签为红色
ax1.tick_params(axis='y', labelcolor='red') # 左轴刻度设为红色

# 5. 创建第二个坐标轴（共享同一个横坐标）
ax2 = ax1.twinx()
# 绘制右轴数据：等概率曲线 (蓝色，菱形标记)
line_prob, = ax2.plot(K, P_equal, 'D-', color='blue', linewidth=2, markersize=8, label='p (equal for all classes)')
ax2.set_ylabel('probability (equal for all classes)', color='blue', fontsize=12) # 右轴标签为蓝色
ax2.tick_params(axis='y', labelcolor='blue') # 右轴刻度设为蓝色

# 6. 设置坐标轴范围，使图形美观
ax1.set_xlim(0, 11) # 横轴稍微留空
ax1.set_ylim(bottom=0) # 熵从0开始
ax2.set_ylim(0, 1.1) # 概率从0到1.1，留出顶部空间

# 7. 添加图例（组合两条曲线的图例）
# 将两条曲线的句柄和标签合并
lines = [line_entropy, line_prob]
labels = [line.get_label() for line in lines]
ax1.legend(lines, labels, loc='upper right', fontsize=10) # 将图例放在左上角

# 8. 添加公式标注（使用LaTeX语法）
# 在图中空白位置添加文本框
formula_text = r'$Entropy(N) = -\sum P(w_j) \log_2 P(w_j)$'
plt.figtext(0.5, 0.01, formula_text, ha='center', fontsize=14, style='italic') # 放在图底部中央

# 9. 添加网格，方便读数
ax1.grid(True, linestyle='--', alpha=0.7)

# 10. 调整布局并显示
plt.title('Relationship between Number of Classes, Max Entropy, and Probability', fontsize=13, pad=20)
fig.tight_layout() # 自动调整子图参数，使之填充整个图像区域
plt.show()

# 11. 保存图片（取消注释下一行来保存）
# fig.savefig('Entropy_vs_Classes.png', dpi=300, bbox_inches='tight')